import base64
import os

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import serialization
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_v1_5

from apps.base.logMsg import logMsg


def create_rsa():
    # 生成一个新的RSA密钥对
    private_key = rsa.generate_private_key(
        backend=default_backend(),
        public_exponent=65537,
        key_size=2048
    )

    # 获取公钥
    public_key = private_key.public_key()

    # 序列化私钥和公钥
    private_pem = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption()
    )
    public_pem = public_key.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo
    )

    # 将私钥保存到文件
    with open('private_key.pem', 'wb') as private_file:
        private_file.write(private_pem)

    # 将公钥保存到文件
    with open('public_key.pem', 'wb') as public_file:
        public_file.write(public_pem)


def encrypto_data(data):
    cipher_text = ''
    with open(os.path.join(os.getcwd(), "public_key.pem"), "r") as f:
        key = f.read()
        rsakey = RSA.importKey(key)
        cipher = PKCS1_v1_5.new(rsakey)
        cipher_text = base64.b64encode(cipher.encrypt(
            data.encode('utf-8')))

    return cipher_text.decode('utf-8')


def decripta_data(data):
    with open(os.path.join(os.getcwd(), "private_key.pem"), "r") as f:
        priv_key = f.read()
    rsa_priv_key = RSA.import_key(priv_key)
    cipher_rsa = PKCS1_v1_5.new(rsa_priv_key)
    keywords = str(data).replace(' ', '+')  # 处理空格或加号的替换
    decoded_data = base64.b64decode(keywords)  # base64 解码
    decrypted_data = cipher_rsa.decrypt(decoded_data, None).decode('utf-8')
    logMsg().info("Decrypted data: {}".format(decrypted_data))
    return decrypted_data
